import sys
import pandas as pd
import talib
import efinance as ef
import numpy as np



def calculate_indicators(df):
    df['日期'] = pd.to_datetime(df['日期'])
    df.set_index('日期', inplace=True)
    df.sort_index(inplace=True)

# ======================================
# 1. 计算 MACD 指标
# ======================================
    macd, signal, hist = talib.MACD(
        df['收盘'].values,
        fastperiod=12,   # 快速EMA周期
        slowperiod=26,   # 慢速EMA周期
        signalperiod=9   # Signal线周期
    )
    df['MACD'] = macd
    df['MACDSignal'] = signal
    df['MACDHist'] = hist

# ======================================
# 2. 计算 RSI 指标 (默认周期为14)
# ======================================
    rsi_14 = talib.RSI(df['收盘'], timeperiod=14)
    df['RSI'] = rsi_14

    # ======================================
    # 3. 计算 KDJ 指标（随机指标 STOCH 的变体）
    # ======================================
    # 使用 talib.STOCH 得到原始的 fast_k 和 slow_d（即 K 和 D）
    k, d = talib.STOCH(
        df['最高'].values,
        df['最低'].values,
        df['收盘'].values,
        fastk_period=9,       # K 周期
        slowk_period=3,       # 平滑 K（相当于 J 的计算基础）
        slowk_matype=0,       # 移动平均类型（0=简单移动平均）
        slowd_period=3,       # D 周期
        slowd_matype=0        # D 的移动平均类型
    )

    df['K'] = k
    df['D'] = d
    df['J'] = 3 * df['K'] - 2 * df['D']
    
    return df.reset_index()



# 修改主函数
if __name__ == '__main__':
    stock_code = sys.argv[1]
    start_date = sys.argv[2] if len(sys.argv) > 2 else None
    end_date = sys.argv[3] if len(sys.argv) > 3 else None
    # stock_code = '603583'
    # start_date = '20250301'
    # end_date = '20250519'

    # session = requests.Session()
    # session.verify = False
    # ef.utils.requests_session = session
# 获取行情数据
    df_quote = ef.stock.get_quote_history(stock_code, beg=start_date, end=end_date, klt=101)

    # 获取资金流向数据（全量）
    df_bill = ef.stock.get_history_bill(stock_code)

    # 转换日期格式
    df_quote['日期'] = pd.to_datetime(df_quote['日期'])
    df_bill['日期'] = pd.to_datetime(df_bill['日期'])

    # 设置索引用于合并
    df_quote.set_index('日期', inplace=True)
    df_bill.set_index('日期', inplace=True)

    # 合并两个 DataFrame（inner join）
    df_merged = pd.merge(df_quote, df_bill, how='inner', left_index=True, right_index=True)

    # 删除重复字段
    df_merged = df_merged.loc[:, ~df_merged.columns.duplicated()]
    df_merged.reset_index(inplace=True)

    # 手动去重：只保留一个“股票名称”、“股票代码”
    df_merged.rename(columns={
        '股票名称_x': '股票名称',
        '股票代码_x': '股票代码',
        '涨跌幅_x': '涨跌幅'
    }, inplace=True)

    # 删除多余的 _y 列
    df_merged.drop(columns=[
        '股票名称_y', '股票代码_y','涨跌幅_y'
    ], errors='ignore', inplace=True)

    # 计算指标
    df_final = calculate_indicators(df_merged)

    # 将日期格式化为 "YYYY-MM-DD"
    df_final['日期'] = df_final['日期'].dt.strftime('%Y-%m-%d')
    print(df_final.to_json(orient='records', force_ascii=False))